import gradio as gr
import numpy as np
from PIL import Image
from pinecone import Pinecone
from collections import Counter

# 初始化 Pinecone
pinecone = Pinecone(api_key="972121d6-7afd-4634-bc4a-3708ce15caaa")
index_name = "mnist-index"

# 连接到 Pinecone 索引
index = pinecone.Index(index_name)

# 预处理图像
def preprocess(image):
    # 将输入的NumPy数组转换为PIL图像对象
    image = Image.fromarray(image)
    image = image.resize((8, 8)).convert('L')
    image_array = np.array(image)
    img_array = image_array.ravel()
    return img_array

# 定义预测函数
def model_predict(img_array):
    print(type(img_array))  # 打印 img_array 的类型
    if isinstance(img_array, np.ndarray):
        img_list = img_array.tolist()
    else:
        raise ValueError("img_array is not a NumPy array")

    # 使用准备好的查询向量在 Pinecone 索引中执行搜索
    results = index.query(
        vector=img_list,
        top_k=11,
        include_metadata=True
    )

    # 从搜索结果中提取每个匹配项的标签
    labels = [match['metadata']['label'] for match in results['matches']]

    # 使用投票机制确定最终的分类结果
    final_prediction = Counter(labels).most_common(1)[0][0]
    return final_prediction

# 创建 Gradio 界面
iface = gr.Interface(
    fn=model_predict,
    inputs=gr.Sketchpad(label="Image", type="numpy"),
    outputs=gr.Label(num_top_classes=1),
    title="knn手写数字识别",
    description="knn预测手写数字"
)

iface.launch(share=True)